from ultralytics import YOLO
import numpy as np

class Detector:
    def __init__(self, model_path):
        """
        初始化检测器。
        :param model_path: YOLO模型的路径。
        """
        self.model = YOLO(model_path)

    def detect(self, frame):
        """
        在给定的帧中检测目标。
        :param frame: 要处理的图像帧 (numpy array)。
        :return: 检测到的目标的边界框列表。
        """
        # 使用YOLO模型进行预测
        results = self.model(frame, classes=[0], verbose=False)  # classes=[0] 表示只检测'person'类别

        # 提取边界框、置信度和类别ID
        boxes = results[0].boxes.xyxy.cpu().numpy()  # Bounding boxes in xyxy format
        confidences = results[0].boxes.conf.cpu().numpy() # Confidence scores
        
        # 将检测结果合并为一个numpy数组 [x1, y1, x2, y2, confidence]
        if len(boxes) > 0:
            detections = np.hstack((boxes, confidences[:, np.newaxis]))
            return detections
        else:
            return np.empty((0, 5))